# -*- coding: utf-8 -*-
#  @author  Bink
#  @date  2020/11/29 15:07
# @Email : 2641032316@qq.com
import tensorflow as tf
from tensorflow.python import keras
from tensorflow.python.keras import layers, models


def h_sigmoid(x):
    output = tf.keras.layers.Activation('hard_sigmoid')(x)
    return output


def h_swish(x):
    output = x * h_sigmoid(x)

    return output


def Squeeze_excitation_layer(x):
    inputs = x
    squeeze = inputs.shape[-1] // 2
    excitation = inputs.shape[-1]
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    x = tf.keras.layers.Dense(squeeze)(x)
    x = tf.keras.layers.Activation('relu')(x)
    x = tf.keras.layers.Dense(excitation)(x)
    x = h_sigmoid(x)
    x = tf.keras.layers.Reshape((1, 1, excitation))(x)
    x = inputs * x

    return x


def BottleNeck(inputs, exp_size, out_size, kernel_size, strides, is_se_existing, activation):
    x = tf.keras.layers.Conv2D(filters=exp_size,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(inputs)
    x = tf.keras.layers.BatchNormalization()(x)
    if activation == "HS":
        x = h_swish(x)
    elif activation == "RE":
        x = tf.keras.layers.Activation(tf.nn.relu6)(x)
    x = tf.keras.layers.DepthwiseConv2D(kernel_size=kernel_size,
                                        strides=strides,
                                        padding="same")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    if activation == "HS":
        x = h_swish(x)
    elif activation == "RE":
        x = tf.keras.layers.Activation(tf.nn.relu6)(x)
    if is_se_existing:
        x = Squeeze_excitation_layer(x)
    x = tf.keras.layers.Conv2D(filters=out_size,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = tf.keras.layers.Activation(tf.keras.activations.linear)(x)
    if strides == 1 and inputs.shape[-1] == out_size:
        x = tf.keras.layers.add([x, inputs])
    return x


def MobileNetV3Large(inputs, classes=10):
    x = tf.keras.layers.Conv2D(filters=16,
                               kernel_size=(3, 3),
                               strides=2,
                               padding="same")(inputs)
    x = tf.cast(x, tf.float32)
    x = tf.keras.layers.BatchNormalization()(x)
    x = h_swish(x)

    x = BottleNeck(x, exp_size=16, out_size=16, kernel_size=3, strides=1, is_se_existing=False, activation="RE")
    x = BottleNeck(x, exp_size=64, out_size=24, kernel_size=3, strides=2, is_se_existing=False, activation="RE")
    x = BottleNeck(x, exp_size=72, out_size=24, kernel_size=3, strides=1, is_se_existing=False, activation="RE")
    x = BottleNeck(x, exp_size=72, out_size=40, kernel_size=5, strides=2, is_se_existing=True, activation="RE")
    x = BottleNeck(x, exp_size=120, out_size=40, kernel_size=5, strides=1, is_se_existing=True, activation="RE")
    x = BottleNeck(x, exp_size=120, out_size=40, kernel_size=5, strides=1, is_se_existing=True, activation="RE")
    x = BottleNeck(x, exp_size=240, out_size=80, kernel_size=3, strides=2, is_se_existing=False, activation="HS")
    x = BottleNeck(x, exp_size=200, out_size=80, kernel_size=3, strides=1, is_se_existing=False, activation="HS")
    x = BottleNeck(x, exp_size=184, out_size=80, kernel_size=3, strides=1, is_se_existing=False, activation="HS")
    x = BottleNeck(x, exp_size=184, out_size=80, kernel_size=3, strides=1, is_se_existing=False, activation="HS")
    x = BottleNeck(x, exp_size=480, out_size=112, kernel_size=3, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=672, out_size=112, kernel_size=3, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=672, out_size=160, kernel_size=5, strides=2, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=960, out_size=160, kernel_size=5, strides=1, is_se_existing=True, activation="HS")
    x = BottleNeck(x, exp_size=960, out_size=160, kernel_size=5, strides=1, is_se_existing=True, activation="HS")

    x = tf.keras.layers.Conv2D(filters=960,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(x)
    x = tf.keras.layers.BatchNormalization()(x)
    x = h_swish(x)
    x = tf.keras.layers.AveragePooling2D(pool_size=(7, 7), strides=1)(x)
    x = tf.keras.layers.Conv2D(filters=1280,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same")(x)
    x = h_swish(x)
    x = tf.keras.layers.Conv2D(filters=classes,
                               kernel_size=(1, 1),
                               strides=1,
                               padding="same",
                               activation=tf.keras.activations.softmax)(x)
    return x


inputs = layers.Input(shape=[224, 244, 3])
outputs = MobileNetV3Large(inputs)

model = models.Model(inputs=inputs, outputs=outputs)

model.compile(optimizer=keras.optimizers.Adam(),
              loss=tf.keras.losses.sparse_categorical_crossentropy,
              metrics=['accuracy'])

model.summary()

(train, train_label), (test, test_label) = keras.datasets.cifar10.load_data()
train, test = train / 255.0, test / 255.0
train = tf.keras.backend.resize_images(train, 7, 7, "channels_last", "nearest")
test = tf.keras.backend.resize_images(test, 7, 7, "channels_last", "nearest")
train_label = tf.convert_to_tensor(train_label)
test_label = tf.convert_to_tensor(test_label)
# print(train.shape)
# print(type(train))
# print(train_label.shape)
# print(test.shape)
# print(test_label.shape)

model.fit(x=train, y=train_label, epochs=1, batch_size=128, validation_data=(test, test_label), validation_freq=1, steps_per_epoch= 1)
